import os
import unittest

import torch
from torch import hub
from torch.backends import cudnn
from torchsummary import summary
from config import Config

from training.models.cbam_convnext_train import cbam_convnext_train

if torch.cuda.is_available() == False:
    # CPU:
    CONFIG = Config()
    hub.set_dir(CONFIG['WINDOWS_TORCH_HOME'])
else:
    # GPU:
    CONFIG = Config()
    hub.set_dir(CONFIG['TORCH_HOME'])
    os.environ["CUDA_VISIBLE_DEVICES"] = CONFIG['CUDA_VISIBLE_DEVICES']
    torch.backends.cudnn.benchmark = True

class CBAM_ResnetTextCase(unittest.TestCase):

    def test_summary_convnext_train(self):

        # model = resnet50(pretrained=False)
        # model = convnext_large_in22k(pretrained=False)
        model = cbam_convnext_train(pretrained=False)
        if torch.cuda.is_available() == False:
            model = model.cpu()
        else:
            model = model.cuda()
        input_size = model.default_cfg['input_size']
        summary(model, input_size=input_size)


if __name__ == '__main__':
    unittest.main()